import threading

from src.utils.error import DeadlockException


class LockManager:
    def __init__(self):
        self.x2u = {}  # 事务持有的资源
        self.u2x = {}  # 资源被哪个事务持有
        self.wait = {}  # 等待队列
        self.wait_locks = {}  # 等待锁
        self.wait_u = {}  # 事务等待的资源
        self.lock = threading.RLock()

    def acquire(self, xid, uid):
        """获取资源锁，可能阻塞或检测死锁"""
        with self.lock:
            # 检查是否已持有
            if xid in self.x2u and uid in self.x2u[xid]:
                return True

            # 检查资源是否可用
            if uid not in self.u2x:
                self.u2x[uid] = xid
                if xid not in self.x2u:
                    self.x2u[xid] = set()
                self.x2u[xid].add(uid)
                return True

            # 需要等待
            self.wait_u[xid] = uid
            if uid not in self.wait:
                self.wait[uid] = []
            self.wait[uid].append(xid)

            # 死锁检测
            if self._has_deadlock():
                # 清理等待状态
                self.wait[uid].remove(xid)
                del self.wait_u[xid]
                raise DeadlockException(f"Deadlock detected for xid {xid}")

            # 创建等待锁
            wait_lock = threading.Event()
            self.wait_locks[xid] = wait_lock

        # 释放全局锁后等待
        wait_lock.wait()
        return True

    def release(self, xid):
        """释放事务持有的所有资源"""
        with self.lock:
            if xid not in self.x2u:
                return

            for uid in list(self.x2u[xid]):
                self._release_one(xid, uid)

            del self.x2u[xid]
    
    def _release_one(self, xid, uid):
        """释放单个资源并选择新的事务来占用该资源"""
        if xid not in self.x2u or uid not in self.x2u[xid]:
            return
        
        self.x2u[xid].remove(uid)
        self._select_new_xid(uid)
    
    def _select_new_xid(self, uid):
        """从等待队列中选择一个事务来占用资源"""
        self.u2x.pop(uid, None)
        wait_list = self.wait.get(uid)
        if wait_list is None or len(wait_list) == 0:
            return
        
        while len(wait_list) > 0:
            xid = wait_list.pop(0)
            if xid not in self.wait_locks:
                continue
            
            # 选择新的事务占用资源
            self.u2x[uid] = xid
            wait_lock = self.wait_locks.pop(xid)
            self.wait_u.pop(xid, None)
            
            # 唤醒等待的事务
            wait_lock.set()
            break
        
        if len(wait_list) == 0:
            self.wait.pop(uid, None)

    def _has_deadlock(self):
        """使用DFS检测等待图中是否有环"""
        visited = {}
        timestamp = 1

        for xid in self.x2u:
            if xid not in visited:
                timestamp += 1
                if self._dfs(xid, visited, timestamp):
                    return True
        return False

    def _dfs(self, xid, visited, timestamp):
        if xid in visited:
            if visited[xid] == timestamp:
                return True  # 发现环
            return False

        visited[xid] = timestamp

        # 检查此事务是否在等待其他事务持有的资源
        if xid not in self.wait_u:
            return False

        uid = self.wait_u[xid]
        if uid not in self.u2x:
            return False

        next_xid = self.u2x[uid]
        return self._dfs(next_xid, visited, timestamp)